import torch
import torch.nn as nn
import torch.nn.functional as F

# 问题 1：为什么使用 LSTM 模型进行词性标注？
# 答案：词性标注依赖单词的上下文信息（如 “read” 在不同句子中可能是动词或名词），LSTM 通过门控机制能够捕捉长距离序列依赖，有效建模上下文关系，适合处理序列标注任务。

# 问题 2：训练过程中为什么需要初始化隐藏层数据？
# 答案：LSTM 的隐藏状态会保存前一次输入的历史信息，若不初始化，处理新句子时会引入无关的历史信息，导致预测偏差。
# 因此，每处理一个新句子前需通过model.init_hidden()将隐藏状态重置为全零张量，确保模型仅基于当前句子进行预测。

# 问题 3：如何将文本数据转换为模型可处理的数值输入？
# 答案：
# 构建索引字典：word_to_ix将每个单词映射到唯一整数索引（如 "The"→0），tag_to_ix将标签映射到索引（如 “DET”→0）。
# 数据预处理函数：prepare_sequence根据索引字典，将单词列表转换为整数张量（如["The", "cat"]→tensor([0, 1])），作为模型的输入数据。



# ======================
# 1. 数据准备
# ======================
# 定义训练数据：每个样本为（句子单词列表， 词性标签列表）
# 词性标签说明：DET=限定词, NN=名词, V=动词
training_data = [
    ("The cat ate the fish".split(), ["DET", "NN", "V", "DET", "NN"]),
    ("They read that book".split(), ["NN", "V", "DET", "NN"])
]

# 定义测试数据：仅包含句子（无标签，用于模型预测）
testing_data = [("They ate the fish".split())]

# 构建单词到索引的映射（词汇表）
word_to_ix = {}
for sentence, tags in training_data:
    for word in sentence:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)  # 自动分配唯一索引
print("单词索引映射:", word_to_ix)  # 输出: {'The': 0, 'cat': 1, 'ate': 2, 'the': 3, 'fish': 4, 'They': 5, 'read': 6, 'that': 7, 'book': 8}

# 手动定义词性标签到索引的映射（标签集）
tag_to_ix = {"DET": 0, "NN": 1, "V": 2}  # 共3种标签


# ======================
# 2. 模型定义（LSTM词性标注器）
# ======================
class LSTMTagger(nn.Module):
    def __init__(self, embedding_dim, hidden_dim, vocab_size, tagset_size):
        super(LSTMTagger, self).__init__()
        self.hidden_dim = hidden_dim  # LSTM隐藏层维度

        # 词嵌入层：将单词索引转换为向量（输入层）
        self.word_embeddings = nn.Embedding(vocab_size, embedding_dim)

        # LSTM层：处理序列数据，输入维度=词嵌入维度，隐藏层维度=hidden_dim
        self.lstm = nn.LSTM(embedding_dim, hidden_dim)

        # 线性层：将LSTM输出映射到标签空间（输出层）
        self.hidden2tag = nn.Linear(hidden_dim, tagset_size)

        # 初始化隐藏状态（用于LSTM的初始输入）
        self.hidden = self.init_hidden()

    def init_hidden(self):
        """初始化LSTM的隐藏状态和细胞状态（全零张量）"""
        return (torch.zeros(1, 1, self.hidden_dim),  # 隐藏状态
                torch.zeros(1, 1, self.hidden_dim))  # 细胞状态

    def forward(self, sentence):
        """前向传播函数"""
        # 1. 词嵌入：将单词索引转换为向量（形状：[句子长度, 嵌入维度]）
        embeds = self.word_embeddings(sentence)

        # 2. LSTM处理：输入需为(序列长度, 批量大小, 特征维度)
        #   这里批量大小为1，所以reshape为(len(sentence), 1, -1)
        lstm_out, self.hidden = self.lstm(embeds.view(len(sentence), 1, -1), self.hidden)

        # 3. 线性变换：将LSTM输出映射到标签分数（形状：[句子长度, 标签数量]）
        tag_space = self.hidden2tag(lstm_out.view(len(sentence), -1))

        # 4. 计算标签概率分布（对数软最大值，便于后续NLLLoss计算）
        tag_scores = F.log_softmax(tag_space, dim=1)
        return tag_scores


# ======================
# 3. 模型初始化与配置
# ======================
# 超参数设置
EMBEDDING_DIM = 6    # 词嵌入向量维度
HIDDEN_DIM = 6       # LSTM隐藏层维度

# 实例化模型（词汇表大小=word_to_ix的长度，标签集大小=tag_to_ix的长度）
model = LSTMTagger(EMBEDDING_DIM, HIDDEN_DIM, len(word_to_ix), len(tag_to_ix))

# 定义损失函数和优化器
loss_function = nn.NLLLoss()  # 负对数似然损失（适用于多分类）
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)  # 随机梯度下降优化器


# ======================
# 4. 数据预处理函数
# ======================
def prepare_sequence(seq, to_ix):
    """将单词/标签列表转换为模型输入的张量（索引序列）"""
    idxs = [to_ix[word] for word in seq]
    return torch.tensor(idxs, dtype=torch.long)


# ======================
# 5. 模型训练
# ======================
for epoch in range(400):  # 训练400轮
    for sentence, tags in training_data:
        # 梯度清零（避免梯度累加）
        model.zero_grad()

        # 初始化LSTM隐藏状态（每轮句子处理前重置）
        model.hidden = model.init_hidden()

        # 数据预处理：转换为索引张量
        sentence_tensor = prepare_sequence(sentence, word_to_ix)
        tags_tensor = prepare_sequence(tags, tag_to_ix)

        # 前向传播：获取标签分数
        tag_scores = model(sentence_tensor)

        # 计算损失：比较预测分数与真实标签
        loss = loss_function(tag_scores, tags_tensor)

        # 反向传播：计算梯度
        loss.backward()

        # 参数更新：优化器根据梯度调整模型参数
        optimizer.step()

    # 每50轮打印一次训练进度（可选）
    if epoch % 50 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")


# ======================
# 6. 模型预测（测试）
# ======================
def predict_tags(sentence):
    """对输入句子进行词性标注预测"""
    # 数据预处理
    sentence_tensor = prepare_sequence(sentence, word_to_ix)
    # 前向传播
    with torch.no_grad():  # 预测时关闭梯度计算
        tag_scores = model(sentence_tensor)
    # 获取每个位置的最大分数标签索引
    _, predicted_indices = torch.max(tag_scores, 1)
    # 将索引映射回标签名称
    predicted_tags = [list(tag_to_ix.keys())[idx] for idx in predicted_indices]
    return predicted_tags

# 对测试数据进行预测
print("\n=== 测试数据预测 ===")
for test_sentence in testing_data:
    print("输入句子:", test_sentence)
    predicted = predict_tags(test_sentence)
    print("预测标签:", predicted)

# 对训练数据进行预测（查看模型学习效果）
print("\n=== 训练数据预测 ===")
for (train_sentence, true_tags) in training_data:
    print("输入句子:", train_sentence)
    print("真实标签:", true_tags)
    predicted = predict_tags(train_sentence)
    print("预测标签:", predicted)
    print("-" * 30)